#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from collections import defaultdict
import logging
from datetime import datetime

from .vpn_knowledge import VPNProtocol, ProtocolFeature, get_protocol_features
from .behavior_analyzer import BehaviorAnalyzer

@dataclass
class DetectionResult:
    """VPN检测结果"""
    protocol: Optional[VPNProtocol]
    confidence: float
    features: Dict[str, float]
    evidence: List[str]

class RAGVPNDetector:
    """基于RAG的VPN检测器"""
    
    def __init__(self, window_size: int = 100):
        self.logger = logging.getLogger(__name__)
        self.behavior_analyzer = BehaviorAnalyzer(window_size)
        self.flow_cache = defaultdict(lambda: {
            'protocol_candidates': set(),
            'confidence_scores': {},
            'evidence': [],
            'last_update': None
        })
        
    def detect(self, packet: Packet) -> DetectionResult:
        """
        检测数据包是否属于VPN流量
        """
        # 1. 提取基础特征
        flow_id = self._get_flow_id(packet)
        behavior_features = self.behavior_analyzer.analyze_behavior(packet, {})
        
        # 2. 基于知识库的初步识别
        candidates = self._identify_candidates(packet, behavior_features)
        
        # 3. 更新流缓存
        self._update_flow_cache(flow_id, candidates, behavior_features)
        
        # 4. 综合判断
        final_protocol, confidence, evidence = self._make_final_decision(flow_id)
        
        return DetectionResult(
            protocol=final_protocol,
            confidence=confidence,
            features=behavior_features,
            evidence=evidence
        )
    
    def _identify_candidates(self, packet: Packet, features: Dict[str, float]) -> List[Tuple[VPNProtocol, float, str]]:
        """基于知识库识别可能的VPN协议"""
        candidates = []
        
        # 1. 基于端口的识别
        if packet.haslayer(TCP):
            port = packet[TCP].dport
            for protocol in get_protocol_by_port(port):
                candidates.append((protocol, 0.6, f"匹配端口 {port}"))
                
        # 2. 基于TLS指纹的识别
        if packet.haslayer(TLS):
            tls_fp = self._extract_tls_fingerprint(packet)
            for protocol in get_protocol_by_tls_fingerprint(tls_fp):
                candidates.append((protocol, 0.7, "匹配TLS指纹"))
                
        # 3. 基于包模式的识别
        for protocol in get_all_protocols():
            feature = get_protocol_features(protocol)
            if self._match_packet_patterns(packet, feature.packet_patterns):
                candidates.append((protocol, 0.8, "匹配包模式"))
                
        # 4. 基于流量特征的识别
        for protocol in get_all_protocols():
            feature = get_protocol_features(protocol)
            if self._match_flow_characteristics(features, feature.flow_characteristics):
                candidates.append((protocol, 0.5, "匹配流量特征"))
                
        return candidates
    
    def _update_flow_cache(self, flow_id: str, candidates: List[Tuple[VPNProtocol, float, str]], features: Dict[str, float]):
        """更新流缓存"""
        cache = self.flow_cache[flow_id]
        current_time = datetime.now()
        
        # 更新候选协议
        for protocol, confidence, evidence in candidates:
            cache['protocol_candidates'].add(protocol)
            if protocol not in cache['confidence_scores']:
                cache['confidence_scores'][protocol] = []
            cache['confidence_scores'][protocol].append(confidence)
            cache['evidence'].append(evidence)
            
        # 清理过期数据
        if cache['last_update'] and (current_time - cache['last_update']).seconds > 300:
            cache['protocol_candidates'].clear()
            cache['confidence_scores'].clear()
            cache['evidence'].clear()
            
        cache['last_update'] = current_time
        
    def _make_final_decision(self, flow_id: str) -> Tuple[Optional[VPNProtocol], float, List[str]]:
        """做出最终判断"""
        cache = self.flow_cache[flow_id]
        
        if not cache['protocol_candidates']:
            return None, 0.0, []
            
        # 计算每个协议的最终置信度
        final_scores = {}
        for protocol in cache['protocol_candidates']:
            scores = cache['confidence_scores'][protocol]
            # 使用加权平均，最近的数据权重更高
            weights = np.linspace(0.5, 1.0, len(scores))
            final_scores[protocol] = np.average(scores, weights=weights)
            
        # 选择置信度最高的协议
        if final_scores:
            best_protocol = max(final_scores.items(), key=lambda x: x[1])
            return best_protocol[0], best_protocol[1], cache['evidence']
            
        return None, 0.0, []
        
    def _extract_tls_fingerprint(self, packet: Packet) -> Dict:
        """提取TLS指纹"""
        fingerprint = {}
        if packet.haslayer(TLS):
            tls = packet[TLS]
            if hasattr(tls, 'cipher_suites'):
                fingerprint['cipher_suites'] = tls.cipher_suites
            if hasattr(tls, 'extensions'):
                fingerprint['extensions'] = tls.extensions
        return fingerprint
        
    def _match_packet_patterns(self, packet: Packet, patterns: List[Dict]) -> bool:
        """匹配包模式"""
        for pattern in patterns:
            layer = pattern['layer']
            if not packet.haslayer(layer):
                continue
                
            layer_packet = packet[layer]
            matches = True
            for key, value in pattern.items():
                if key != 'layer' and not hasattr(layer_packet, key):
                    matches = False
                    break
                if key != 'layer' and getattr(layer_packet, key) != value:
                    matches = False
                    break
                    
            if matches:
                return True
                
        return False
        
    def _match_flow_characteristics(self, features: Dict[str, float], characteristics: Dict) -> bool:
        """匹配流量特征"""
        if 'min_packet_size' in characteristics:
            if features.get('avg_packet_size', 0) < characteristics['min_packet_size']:
                return False
                
        if 'max_packet_size' in characteristics:
            if features.get('avg_packet_size', 0) > characteristics['max_packet_size']:
                return False
                
        if 'protocol_ratio' in characteristics:
            for proto, ratio in characteristics['protocol_ratio'].items():
                if features.get(f'{proto.lower()}_ratio', 0) < ratio * 0.8:
                    return False
                    
        return True 